use image;

// 翻译宏定义的变量
const PTP_ENTRIES: usize = 512;
const PTP_SIZE: usize = 4096;
const SIZE_2M: usize = 2 * 1024 * 1024;
const IS_VALID: u64 = 1 << 0;
const IS_TABLE: u64 = 1 << 1;
const UXN: u64 = 1 << 54;
const ACCESSED: u64 = 1 << 10;
const INNER_SHARABLE: u64 = 3 << 8;
const NORMAL_MEMORY: u64 = 4 << 2;
const DEVICE_MEMORY: u64 = 0 << 2;
// 使用了repr(align(PTP_SIZE))来确保PageTableLevel
// 和其嵌套的类型在内存中对齐到PTP_SIZE，
// 与C代码中的ALIGN宏相似。https://rustwiki.org/zh-CN/edition-guide/rust-2018/data-types/choosing-alignment-with-the-repr-attribute.html
#[repr(align(PTP_SIZE))]
struct PageTableLevel {
    entries: [u64; PTP_ENTRIES],
}

static mut BOOT_TTBR0_L0: PageTableLevel = PageTableLevel { entries: [0; PTP_ENTRIES] };
static mut BOOT_TTBR0_L1: PageTableLevel = PageTableLevel { entries: [0; PTP_ENTRIES] };
static mut BOOT_TTBR0_L2: PageTableLevel = PageTableLevel { entries: [0; PTP_ENTRIES] };

static mut BOOT_TTBR1_L0: PageTableLevel = PageTableLevel { entries: [0; PTP_ENTRIES] };
static mut BOOT_TTBR1_L1: PageTableLevel = PageTableLevel { entries: [0; PTP_ENTRIES] };
static mut BOOT_TTBR1_L2: PageTableLevel = PageTableLevel { entries: [0; PTP_ENTRIES] };

const PHYSMEM_START: usize = 0x0;
const PHYSMEM_BOOT_END: usize = 0x10000000;
const PERIPHERAL_BASE: usize = 0x20000000;
const PHYSMEM_END: usize = 0x40000000;

const KERNEL_VADDR: usize = 0xFFFF_FFFF_8000_0000;

// 直接使用普通函数来代替C中的宏定义函数
fn get_l0_index(x: usize) -> usize {
    (x >> (12 + 9 + 9 + 9)) & 0x1ff
}

fn get_l1_index(x: usize) -> usize {
    (x >> (12 + 9 + 9)) & 0x1ff
}

fn get_l2_index(x: usize) -> usize {
    (x >> (12 + 9)) & 0x1ff
}

fn init_boot_pt(){
    // TTBR0_EL1 0-1G
    unsafe {
        BOOT_TTBR0_L0.entries[0] = (&BOOT_TTBR0_L1 as *const _ as u64) | IS_TABLE | IS_VALID;
        BOOT_TTBR0_L1.entries[0] = (&BOOT_TTBR0_L2 as *const _ as u64) | IS_TABLE | IS_VALID;
    }
    // Usable memory: PHYSMEM_START ~ PERIPHERAL_BASE
    let start_entry_idx = PHYSMEM_START / SIZE_2M;
    let end_entry_idx = PERIPHERAL_BASE / SIZE_2M;

    // Map each 2M page
    for idx in start_entry_idx..end_entry_idx {
        unsafe {
            BOOT_TTBR0_L2.entries[idx] = (PHYSMEM_START + idx * SIZE_2M) |
                UXN | ACCESSED | INNER_SHARABLE | NORMAL_MEMORY | IS_VALID;
        }
    }

    // Peripheral memory: PERIPHERAL_BASE ~ PHYSMEM_END

    // raspi3b/3b+ Peripherals: 0x3f 00 00 00 - 0x3f ff ff ff
    let start_entry_idx = end_entry_idx;
    let end_entry_idx = PHYSMEM_END / SIZE_2M;
    // Map each 2M page
    for idx in start_entry_idx..end_entry_idx {
        unsafe {
            BOOT_TTBR0_L2.entries[idx] = (PHYSMEM_START + idx * SIZE_2M) |
                UXN | ACCESSED | DEVICE_MEMORY | IS_VALID;
        }
    }

    // TTBR1_EL1 0-1G
    // KERNEL_VADDR: L0 pte index: 510; L1 pte index: 0; L2 pte index: 0.
    let kva = KERNEL_VADDR;
    unsafe {
        BOOT_TTBR1_L0[GET_L0_INDEX(kva)] = (&BOOT_TTBR1_L1 as *const _ as u64) | IS_TABLE | IS_VALID;
        BOOT_TTBR1_L1[GET_L1_INDEX(kva)] = (&BOOT_TTBR1_L2 as *const _ as u64) | IS_TABLE | IS_VALID;
    }

    let start_entry_idx = GET_L2_INDEX(kva);
    let end_entry_idx = start_entry_idx + PHYSMEM_BOOT_END / SIZE_2M;

    // Map each 2M page
    for idx in start_entry_idx..end_entry_idx {
        unsafe {
            BOOT_TTBR1_L2[idx] = (PHYSMEM_START + idx * SIZE_2M) |
                UXN | ACCESSED | INNER_SHARABLE | NORMAL_MEMORY | IS_VALID;
        }
    }

    // Usable memory: PHYSMEM_START ~ PERIPHERAL_BASE
    let start_entry_idx = start_entry_idx;
    let end_entry_idx = start_entry_idx + PERIPHERAL_BASE / SIZE_2M;

    // Map each 2M page
    for idx in start_entry_idx..end_entry_idx {
        unsafe {
            BOOT_TTBR1_L2[idx] = (PHYSMEM_START + idx * SIZE_2M) |
                UXN | ACCESSED | INNER_SHARABLE | NORMAL_MEMORY | IS_VALID;
        }
    }

    // Peripheral memory: PERIPHERAL_BASE ~ PHYSMEM_END
    let start_entry_idx = start_entry_idx + PERIPHERAL_BASE / SIZE_2M;
    let end_entry_idx = PHYSMEM_END / SIZE_2M;

    // Map each 2M page
    for idx in start_entry_idx..end_entry_idx {
        unsafe {
            BOOT_TTBR1_L2[idx] = (PHYSMEM_START + idx * SIZE_2M) |
                UXN | ACCESSED | DEVICE_MEMORY | IS_VALID;
        }
    }

    // Local peripherals, e.g., ARM timer, IRQs, and mailboxes
    // 0x4000_0000 .. 0xFFFF_FFFF
    // 1G is enough. Map 1G page here.
    let kva = KERNEL_VADDR + PHYSMEM_END;
    unsafe {
        BOOT_TTBR1_L1[GET_L1_INDEX(kva)] = PHYSMEM_END | UXN | ACCESSED | DEVICE_MEMORY | IS_VALID;
    }
}